import os
import pandas as pd
import numpy as np
import json
import random
import glob
from shutil import copyfile, move
from pathlib import Path

# path = r'/root/project/AutoRepair_T7/data/75/all'

# train_path = r'/root/project/AutoRepair_T7/data/75/train'
# test_path = r'/root/project/AutoRepair_T7/data/75/test'
# val_ratio = 0.15

# for code in os.listdir(path):
#     l = glob.glob(os.path.join(path, code, '*.jpg'))
#     if code in ["TGXID", "TCVPT"]:
#         random.shuffle(l)
#         l = l[:100]
#     train_num = int(len(l) * (1-val_ratio))
#     for i in range(train_num):
#         os.makedirs(os.path.join(train_path, code), exist_ok=True)
#         img_name = Path(l[i]).name
#         copyfile(l[i], os.path.join(train_path, code, img_name))
#         if(code != 'TGXID'):
#             json_name = img_name[:-3] + 'json'
#             copyfile(l[i][:-3]+'json', os.path.join(train_path, code, json_name))
    
#     for i in range(train_num, len(l)):
#         os.makedirs(os.path.join(test_path, code), exist_ok=True)
#         img_name = Path(l[i]).name
#         copyfile(l[i], os.path.join(test_path, code, img_name))
#         if(code != 'TGXID'):
#             json_name = img_name[:-3] + 'json'
#             copyfile(l[i][:-3]+'json', os.path.join(test_path, code, json_name))



# path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6006/for_label_0315'

# df = pd.DataFrame() 
# df['imgs'] = glob.glob(os.path.join(path, "*/*.jpg"))
# df['code'] = df['imgs'].apply(lambda x: Path(x).parent.name)

# print(df['code'].value_counts())
# print(len(df))

def move_file(img_path, code, memo, memo2):
    print(img_path)
    print('dst_code: ', code)
    img_name = Path(img_path).name
    for dp in memo.keys():
        if img_name not in memo[dp]:
            continue
        
        for src_code in memo2[dp]:
            if img_name in memo2[dp][src_code]:
                print('src_code: ', src_code)
                if(code == src_code): break
                img_src_path = os.path.join(dp, src_code, img_name)
                img_dst_path = os.path.join(dp, code, img_name)
                os.makedirs(Path(img_dst_path).parent, exist_ok=True)
                move(img_src_path, img_dst_path)

                label_name = str(Path(img_name).with_suffix('.json'))
                label_src_path = os.path.join(dp, src_code, label_name)
                if os.path.exists(label_src_path):
                    if code == 'TGXID':
                        os.remove(label_src_path)
                    else:
                        label_dst_path = os.path.join(dp, code, label_name)
                        move(label_src_path, label_dst_path)
                else:
                    print(f"{label_src_path} not exist! ")
                break
        
        break


src_path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/data_reviewed_0708/T7002'
# dst_paths = [r'/data2/autorepair/ruanzhifeng/tianxinji/data/T6_Vtech/T6009/0621',
#             # r'/data2/autorepair/ruanzhifeng/tianxinji/data/T6_Vtech/T6010/add_data_0705',
#             # r'/data2/autorepair/ruanzhifeng/tianxinji/data/T6_Vtech/T6010/add_data_0712'
            # ]
dst_paths = []
r = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T7002'
for folder in os.listdir(r):
    dst_paths.append(os.path.join(r, folder))

memo = {}
for dp in dst_paths:
    imgs = glob.glob(os.path.join(dp, "**/*.jpg"), recursive=True)
    memo[dp] = [Path(x).name for x in imgs]

memo2 = {}
for dp in dst_paths:
    memo2[dp] = {}
    for code in os.listdir(dp):
        imgs = glob.glob(os.path.join(dp, code, "**/*.jpg"), recursive=True)
        memo2[dp][code] = [Path(x).name for x in imgs]

imgs = glob.glob(os.path.join(src_path, "**/*.jpg"), recursive=True)
df = pd.DataFrame(data={'img_path': imgs})
df['code'] = df['img_path'].apply(lambda x: Path(x).parent.name)

df.apply(lambda x: move_file(img_path=x['img_path'], code=x['code'], memo=memo, memo2=memo2), axis=1)


# # tmp_imgs = glob.glob(os.path.join(dst_path, "*/*.jpg"))
# df = pd.DataFrame() 
# df['imgs'] = glob.glob(os.path.join(dst_path, "*/*.jpg"))
# df['img_name'] = df['imgs'].apply(lambda x: Path(x).name)
# # df['code'] = df['imgs'].apply(lambda x: Path(x).parent.name)
# ex_imgs = set(df['img_name'].tolist())


# df2 = pd.DataFrame() 
# df2['imgs'] = glob.glob(os.path.join(src_path, "*/*.jpg"))
# df2['img_name'] = df['imgs'].apply(lambda x: Path(x).name)
# df2['code'] = df2['imgs'].apply(lambda x: Path(x).parent.name)

# df2['to_del'] = df2['img_name'].apply(lambda x: x in ex_imgs)

# df2 = df2[df2['to_del']]
# df2.apply(lambda x: os.remove(x['imgs']), axis=1)

# # print(df2['to_del'].value_counts())
